import argparse
import os
import sys
import time
from typing import Optional, Tuple, Any

import cv2
import numpy as np

# 解决包导入路径问题（目录名含有连字符，无法作为包名导入）
CURRENT_DIR = os.path.abspath(os.path.dirname(__file__))
RS2_DIR = os.path.abspath(os.path.join(CURRENT_DIR, "..", "rs2"))
if CURRENT_DIR not in sys.path:
    sys.path.append(CURRENT_DIR)
if RS2_DIR not in sys.path:
    sys.path.append(RS2_DIR)

from ZH_yolo_21_opencv import YoloInspector  # noqa: E402
try:
    from GStreamerWrapper import GStreamerWrapper  # type: ignore  # noqa: E402
except Exception:  # noqa: BLE001
    GStreamerWrapper = None  # type: ignore


class CameraBase:
    """
    抽象相机接口。
    子类需实现 open、read、release 三个方法。
    """

    def open(self) -> bool:
        raise NotImplementedError

    def read(self) -> Tuple[bool, Optional[np.ndarray]]:
        raise NotImplementedError

    def release(self) -> None:
        raise NotImplementedError


class OpenCVCamera(CameraBase):
    def __init__(self, device_index: int = 0, width: Optional[int] = None, height: Optional[int] = None) -> None:
        self.device_index = device_index
        self.width = width
        self.height = height
        self.cap: Optional[cv2.VideoCapture] = None

    def open(self) -> bool:
        self.cap = cv2.VideoCapture(self.device_index)
        if not self.cap.isOpened():
            return False
        if self.width is not None:
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
        if self.height is not None:
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
        return True

    def read(self) -> Tuple[bool, Optional[np.ndarray]]:
        if self.cap is None:
            return False, None
        ret, frame = self.cap.read()
        if not ret:
            return False, None
        return True, frame

    def release(self) -> None:
        if self.cap is not None:
            self.cap.release()
            self.cap = None


class GStreamerCamera(CameraBase):
    def __init__(self, width: int = 640, height: int = 360) -> None:
        self.width = width
        self.height = height
        self.wrapper: Any = None

    def open(self) -> bool:
        # GStreamerWrapper 内部会启动拉流和取帧线程
        if GStreamerWrapper is None:
            print("错误：未能导入 GStreamerWrapper，请确认依赖和路径是否正确。")
            return False
        self.wrapper = GStreamerWrapper(width=self.width, height=self.height)
        return True

    def read(self) -> Tuple[bool, Optional[np.ndarray]]:
        if self.wrapper is None:
            return False, None
        frame = self.wrapper.GetFrame()
        if frame is None:
            return False, None
        return True, frame

    def release(self) -> None:
        if self.wrapper is not None:
            try:
                self.wrapper.StopThread()
            except Exception:
                pass
            self.wrapper = None


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="基于抽象相机接口的YOLO巡检（YOLO定位+HSV颜色）：支持OpenCV摄像头与GStreamer拉流"
    )
    parser.add_argument(
        "--camera",
        choices=["cv", "gstreamer"],
        default="cv",
        help="选择相机源：cv为OpenCV本地摄像头，gstreamer为GStreamer拉流",
    )
    parser.add_argument(
        "--device",
        type=int,
        default=0,
        help="当使用--camera=cv时的摄像头索引",
    )
    parser.add_argument("--width", type=int, default=None, help="期望的图像宽度")
    parser.add_argument("--height", type=int, default=None, help="期望的图像高度")
    parser.add_argument(
        "--model",
        type=str,
        default=r"D:\\python.learn\\python.project\\yolo_dataset\\250730\\runs\\0_epochs100\\weights\\best.pt",
        help="YOLO 模型 .pt 文件路径",
    )
    parser.add_argument("--conf", type=float, default=0.6, help="检测置信度阈值")
    return parser


def create_camera(args: argparse.Namespace) -> CameraBase:
    if args.camera == "cv":
        return OpenCVCamera(device_index=args.device, width=args.width, height=args.height)

    # gstreamer: 若未提供分辨率，则使用 GStreamerWrapper 缺省 640x360
    width = args.width if args.width is not None else 640
    height = args.height if args.height is not None else 360
    return GStreamerCamera(width=width, height=height)


def main() -> None:
    parser = build_arg_parser()
    args = parser.parse_args()

    try:
        inspector = YoloInspector(args.model, conf_threshold=args.conf)
    except Exception as exc:
        print(f"错误：无法加载模型，请检查路径 '{args.model}' 是否正确。错误信息: {exc}")
        return

    camera = create_camera(args)
    if not camera.open():
        if args.camera == "cv":
            print("错误：无法打开本地摄像头。")
        else:
            print("错误：无法启动 GStreamer 拉流。")
        return

    print("\n相机已启动。")
    print(
        ">>> 启动混合识别模式：YOLO定位，HSV识别颜色。<<<"
    )
    print("正在等待巡检目标... 按 Ctrl+C 退出。")

    try:
        while True:
            ret, frame = camera.read()
            if not ret or frame is None:
                time.sleep(0.01)
                continue
            inspector.process_frame(frame)
            time.sleep(0.01)

    except KeyboardInterrupt:
        print("\n检测已由用户停止。")
    finally:
        camera.release()
        try:
            pass
        finally:
            cv2.destroyAllWindows()
            print("资源已释放。")


if __name__ == "__main__":
    main()


